package com.example.demo.mcp.webclient.config;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;
import java.util.List;

/**
 * AI配置类
 * 配置MCP客户端和ChatClient
 *
 * @author zwg
 * @date 2025/2/20 22:25
 */
@Slf4j
@Configuration
public class AiConfig {


    @Bean
    public McpSyncClient mcpSyncClientOfAmap() {
        String apiKey = System.getenv("DASHSCOPE_API_KEY");
        // 创建MCP同步客户端实例
        var transport = HttpClientSseClientTransport.builder("https://dashscope.aliyuncs.com")
                .sseEndpoint("/api/v1/mcps/amap-maps/sse")
                .customizeRequest( request-> request.header("Authorization","Bearer " + apiKey))
                .build();
        McpSyncClient mcpSyncClient = McpClient.sync(transport)
                .requestTimeout(Duration.ofSeconds(30))
                .clientInfo(new McpSchema.Implementation("DashScopeAMapMcpClient", "1.0.0"))
                .toolsChangeConsumer((List<McpSchema.Tool> tools) -> {
                    // Handle tools change
                    if(!tools.isEmpty()){
                        log.info("DashScopeAMapMcpClient工具已更新，{} 个工具已可用", tools.size());
                    }
                })
                .loggingConsumer((McpSchema.LoggingMessageNotification logMsg) -> {
                    // Handle log messages
                    if(logMsg.level() == McpSchema.LoggingLevel.DEBUG){
                        log.debug("{}", logMsg.data());
                    }else if(logMsg.level() == McpSchema.LoggingLevel.INFO){
                        log.info("{}", logMsg.data());
                    }else if(logMsg.level() == McpSchema.LoggingLevel.WARNING){
                        log.warn("{}", logMsg.data());
                    }else if(logMsg.level() == McpSchema.LoggingLevel.ERROR){
                        log.error("{}", logMsg.data());
                    }else{
                        log.info("{}", logMsg.data());
                    }
                })
                .build();
        // 初始化客户端
        mcpSyncClient.initialize();
        return mcpSyncClient;
    }

    //@Bean
    public McpSyncClient mcpSyncClientOfLocal() {
        // 创建MCP同步客户端实例
        var transport = HttpClientSseClientTransport.builder("http://localhost:8081")
                .build();
        McpSyncClient mcpSyncClient = McpClient.sync(transport)
                .requestTimeout(Duration.ofSeconds(30))
                .clientInfo(new McpSchema.Implementation("LocalDemoMcpClient", "1.0.0"))
                .build();
        // 初始化客户端
        mcpSyncClient.initialize();
        return mcpSyncClient;
    }

    /**
     * 日志记录器
     */
    @Bean
    public Advisor loggerAdvisor() {
        return new SimpleLoggerAdvisor();
    }


    /**
     * 创建ChatClient实例，集成MCP工具
     */
    @ConditionalOnProperty(name = "spring.ai.model.chat", havingValue = "dashScope")
    @Bean(name = "dashScopeChatClient")
    public ChatClient dashScopeChatClient(DashScopeChatModel dashScopeChatModel, Advisor loggerAdvisor, List<McpSyncClient> mcpClients) {
        SyncMcpToolCallbackProvider mcpToolProvider = SyncMcpToolCallbackProvider.builder().mcpClients(mcpClients).build();
        ToolCallback[] toolCallbacks = mcpToolProvider.getToolCallbacks();
        log.info("创建ChatClient，集成了 {} 个MCP客户端", mcpClients.size());
        // 打印可用的MCP客户端
        for (int i = 0; i < mcpClients.size(); i++) {
            McpSyncClient client = mcpClients.get(i);
            McpSchema.ListToolsResult toolsResult = client.listTools();
            log.info("客户端-{}: name:{} version:{} toolSize:{}",1+i,client.getClientInfo().name(),
                    client.getClientInfo().version(),toolsResult.tools().size());
            // 列出可用工具
            toolsResult.tools().forEach(tool -> {
                log.info("客户端 name:{} tool:{}: {} {}",client.getClientInfo().name(),
                        tool.name() , tool.description(), tool.inputSchema().properties());
            });
        }

        return ChatClient.builder(dashScopeChatModel)
                .defaultAdvisors(loggerAdvisor)
                .defaultToolCallbacks(mcpToolProvider)
                .build();
    }
}
